/* Copyright 2012 Tobias Marschall
 *
 * This file is part of CLEVER.
 *
 * CLEVER is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * CLEVER is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with CLEVER.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <iostream>

#include <boost/program_options.hpp>

#include <bamtools/api/BamReader.h>
#include <bamtools/api/BamWriter.h>

#include "ASTagCalculator.h"
#include "GroupWiseBamReader.h"

using namespace std;
namespace po = boost::program_options;

void usage(const char* name, const po::options_description& options_desc) {
	cerr << "Usage: " << name << " [options] <reference.fasta>" << endl;
	cerr << endl;
	cerr << "Reads BAM format from stdin and outputs BAM format to stdout, where " << endl;
	cerr << "AS tags containing a (re-)computed alignment scores have been added." << endl;
	cerr << "The alignment scores are the sum of base qualities at mismatch positions" << endl;
	cerr << "as extracted from the phred string." << endl;
	cerr << endl;
	cerr << "If given filename ends on \".gz\", its contents is unzipped." << endl;
	cerr << endl;
	cerr << "NOTE: The first word of the sequence names in <reference.fasta> is used," << endl;
	cerr << "i.e. if fasta line is " << endl;
	cerr << ">chr1 xyz" << endl;
	cerr << "then only \"chr1\" is used and searched for in the BAM/SAM file." << endl;
	cerr << endl;
	cerr << options_desc << endl;
	exit(1);
}

int main(int argc, char* argv[]) {
	// PARAMETERS
	int phred_offset;
	int bad_score_threshold;
	bool skip_unknown = false;

	po::options_description options_desc("Allowed options");
	options_desc.add_options()
		("phred_offset,p", po::value<int>(&phred_offset)->default_value(33), "Value to substract from ASCII code to get the PHRED quality.")
		("bad_alignment_threshold,b", po::value<int>(&bad_score_threshold)->default_value(1000), "Issue a warning when AS tag is above this value.")
		("skip_unknown,s", po::value<bool>(&skip_unknown)->zero_tokens(), "Do not output reads for which no AS tag could be computed, e.g. because the reference sequence was unknown.")
	;

	if (isatty(fileno(stdin)) || (argc<2)) {
		usage(argv[0], options_desc);
	}
	string reference_filename(argv[argc-1]);
	argc -= 1;

	po::variables_map options;
	try {
		po::store(po::parse_command_line(argc, argv, options_desc), options);
		po::notify(options);
	} catch(exception& e) {
        cerr << "error: " << e.what() << "\n";
        return 1;
    }

	ASTagCalculator tag_calculator(reference_filename,phred_offset,bad_score_threshold);
	cerr << "Done reading reference sequences" << endl;

	GroupWiseBamReader* bam_reader = new GroupWiseBamReader("/dev/stdin", true);
	bam_reader->enableProgressMessages(cerr, 200000);

	const BamTools::RefVector& bam_ref_data = bam_reader->getReferenceData();
	const BamTools::SamHeader sam_header = bam_reader->getHeader();
	BamTools::BamWriter bam_writer;
	if (!bam_writer.Open("/dev/stdout", sam_header, bam_ref_data)) {
		cerr << "Error writing BAM to stdout." << endl;
		return 1;
	}
	
	while ( bam_reader->hasNext() ) {
		bam_reader->advance();
		
		const vector<BamTools::BamAlignment>& alignments1 = bam_reader->getAlignmentsFirst();
		const vector<BamTools::BamAlignment>& alignments2 = bam_reader->getAlignmentsSecond();	
								
		bool is_primary_reverse = false;
		std::string primary_qualities;
		BamTools::BamAlignment read_aln;
		
		bool primary_found = false; 
		for (size_t i=0; i<alignments1.size(); ++i) {	
		  read_aln = alignments1[i];
		  if (read_aln.IsPrimaryAlignment()) {		    		    	    
		    primary_qualities.assign(read_aln.Qualities.data(), read_aln.Qualities.size());		   
		    is_primary_reverse = read_aln.IsReverseStrand();
		    tag_calculator.computeTag(read_aln, bam_ref_data);
		    bam_writer.SaveAlignment(read_aln);		    		   			    
		    primary_found = true;
		    break;
		  }		  
		}
		if(!primary_found) {
		  cerr << "no primary alignment found." << endl;
		  continue; 
		}		  
		
		std::string rev_primary_qualities (primary_qualities.rbegin(), primary_qualities.rend());				
		
		for(size_t i=0; i<alignments1.size(); ++i) {		  
		  read_aln = alignments1[i];
		  if (!read_aln.IsPrimaryAlignment()) {		    
		    if (read_aln.IsReverseStrand() != is_primary_reverse) {			      		    
		      tag_calculator.computeTag(read_aln, bam_ref_data, rev_primary_qualities);		      
		    }
		    else {		     
		      tag_calculator.computeTag(read_aln, bam_ref_data, primary_qualities);		      
		    }		   		    		    
		    bam_writer.SaveAlignment(read_aln);		   
		  }
		}
		
		is_primary_reverse = false;
		
		primary_found = false;
		for (size_t i=0; i<alignments2.size(); ++i) {	
		  read_aln = alignments2[i];
		  if (read_aln.IsPrimaryAlignment()) {		    		    	    
		    primary_qualities.assign(read_aln.Qualities.data(), read_aln.Qualities.size());		   
		    is_primary_reverse = read_aln.IsReverseStrand();
		    tag_calculator.computeTag(read_aln, bam_ref_data);
		    bam_writer.SaveAlignment(read_aln);		    		   			    
		    primary_found = true;
		    break;
		  }		  
		}
		if(!primary_found) {
		  cerr << "no primary alignment found." << endl;
		  continue; 
		}
		
		rev_primary_qualities = string(primary_qualities.rbegin(), primary_qualities.rend());
		
		for(size_t i=0; i<alignments2.size(); ++i) {		  
		  read_aln = alignments2[i];
		  if (!read_aln.IsPrimaryAlignment()) {		    
		    if (read_aln.IsReverseStrand() != is_primary_reverse) {			      		    
		      tag_calculator.computeTag(read_aln, bam_ref_data, rev_primary_qualities); 		      
		    }
		    else {		     
		      tag_calculator.computeTag(read_aln, bam_ref_data, primary_qualities);		      
		    }		   		    		    
		    bam_writer.SaveAlignment(read_aln);		   
		  }
		}
		
		
	}
	tag_calculator.printWarnings(cerr);
	return 0;
}
